package com.zhangdy.agent.common.utils;

import com.taobao.arthas.compiler.DynamicCompiler;
import com.zhangdy.agent.common.dto.MethodPoint;
import javassist.*;
import org.apache.commons.io.FileUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.util.StringUtils;

import java.io.File;
import java.io.FileInputStream;
import java.util.List;
import java.util.Map;

public class DynamicCompilerUtil {
    private static Logger logger = LoggerFactory.getLogger(DynamicCompilerUtil.class);

    public static final String CUSTOMIZE_CLZ_DIR = "/Users/zhangdaye/zhangdy-agent";
//    private static final String CUSTOMIZE_CLZ_DIR = "E:\\run-jar\\zhangdy-agent";

    public static final String LAUNCHED_CLASS_LOADER = "org.springframework.boot.loader.LaunchedURLClassLoader";

    public static void arthasDynamicCompiler(String className, String sourceCode) throws Exception {
        // 使用 Arthas 动态编译
        DynamicCompiler dynamicCompiler = new DynamicCompiler(Thread.currentThread().getContextClassLoader());
        dynamicCompiler.addSource(className, sourceCode);
        Map<String, byte[]> byteCodes = dynamicCompiler.buildByteCodes();
        File outputDir = new File(CUSTOMIZE_CLZ_DIR);
        for (Map.Entry<String, byte[]> entry : byteCodes.entrySet()) {
            logger.info("dynamic compiler class:{}", entry.getKey());
            File byteCodeFile = new File(outputDir, StringUtils.replace(entry.getKey(), ".", File.separator) + ".class");
            FileUtils.writeByteArrayToFile(byteCodeFile, entry.getValue());
        }
    }

    public static String loadByteCode(String className, String classPath) throws Exception {
        ClassPool classPool = ClassPool.getDefault();
//        ClassPool classPool =  new ClassPool(true);

        CtClass executeClass = null;
        try {
            executeClass = classPool.get(className);
        }catch (Exception e) {
            executeClass = classPool.makeClass(new FileInputStream(classPath));
            ClassLoader contextClassLoader = Thread.currentThread().getContextClassLoader();
            if (!LAUNCHED_CLASS_LOADER.equals(contextClassLoader.getClass().getName())) {
                if (LAUNCHED_CLASS_LOADER.equals(contextClassLoader.getParent().getClass().getName())) {
                    contextClassLoader = contextClassLoader.getParent();
                } else {
                    contextClassLoader = ClassLoader.getSystemClassLoader();
                }
            }
            // 传入类加载器
            executeClass.toClass(contextClassLoader, null);
        }

        //读取字节码文件
//        // 当前上下文类加载器
//        System.out.println("----> current thread context classLoader : " + Thread.currentThread().getContextClassLoader().toString());
//        // 当前上下文类加载器的父类加载器
//        System.out.println("----> current thread context classLoader's parent classLoader : " + Thread.currentThread().getContextClassLoader().getParent().toString());
//        // 应用程序类加载器
//        System.out.println("----> application classLoader : " + ClassLoader.getSystemClassLoader().toString());
//        // 定义 Class 实例\\\

        // 上下文类加载器

        String result = builderExecuteBody(executeClass, false);

        executeClass.defrost();
        return result;


    }

    public static String builderExecuteBody(CtClass executeClass, boolean returnValue) {
        StringBuilder executeBody = new StringBuilder("{").append("\r\n");
        executeBody
                .append(executeClass.getName()) // 类型
                .append(" ")
                .append(executeClass.getSimpleName()) // 变量名称
                .append(" = ")
                .append("new ").append(executeClass.getName()).append("();")
                .append("\r\n");
        // 若是有返回值，则使用临时变量存储
        if (returnValue) {
            executeBody.append("Object result = ");
        }
        executeBody
                .append(executeClass.getSimpleName()).append(".execute($args);")
                .append("\r\n");
        if (returnValue) {
            executeBody.append("return ($r) result;").append("\r\n");
        }
        executeBody.append("}").append("\r\n");
        return executeBody.toString();
    }

    public static CtClass replaceByteCode(MethodPoint point, ClassPool classPool)  throws Exception{
        ClassLoader classLoader = classPool.getClassLoader();

        ClassLoader cl = Thread.currentThread().getContextClassLoader().getParent();
        classPool.appendClassPath(new LoaderClassPath(cl));

        String targetClassName = point.getPackageName() + "." + point.getClassName();
        // 目标类
        CtClass targetClass = classPool.get(targetClassName);
        // 根据参数类型构建 CtClass 数组
        CtClass[] params = buildParams(classPool, point);
        // 目标方法
        CtMethod targetMethod = targetClass.getDeclaredMethod(point.getMethodName(), params);
        // 前置规则
        String beforeCode = point.getBeforeCode();

        // 替换规则
        String replaceCode = point.getReplaceCode();

        // 后置规则
        String afterCode = point.getAfterCode();


        if (replaceCode != null && !"".equals(replaceCode.trim())) {
            // 替换方法
            targetMethod.setBody(replaceCode);
        }

        // 在方法体前插入代码块
        if (beforeCode != null && !"".equals(beforeCode.trim())) {
            targetMethod.insertBefore(beforeCode);
        }

        // 在方法体后插入代码块
        if (afterCode != null && !"".equals(afterCode.trim())) {
            targetMethod.insertAfter(afterCode);
        }

         return targetClass;

    }


    private static CtClass[] buildParams(ClassPool classPool, MethodPoint point) throws Exception {
        CtClass[] result = null;
        List<String> argsTypeList = point.getArgsTypeList();
        if (argsTypeList == null || argsTypeList.isEmpty()) {
            return result;
        }
        result = new CtClass[argsTypeList.size()];
        int i = 0;
        for (String argType : argsTypeList) {
            result[i++] = classPool.get(argType);
        }
        return result;
    }

}
